"""
视频处理模块
负责视频文件的处理、音频提取和格式转换
"""

import os
import logging
from pathlib import Path
from typing import Optional, Dict, Any
from moviepy import VideoFileClip, AudioFileClip
import ffmpeg
import math

class VideoProcessor:
    """视频处理器类"""
    
    SUPPORTED_VIDEO_FORMATS = {'.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.m4v'}
    SUPPORTED_AUDIO_FORMATS = {'.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a'}
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
    def is_supported_format(self, file_path: str) -> bool:
        """检查文件格式是否支持"""
        file_ext = Path(file_path).suffix.lower()
        return file_ext in self.SUPPORTED_VIDEO_FORMATS or file_ext in self.SUPPORTED_AUDIO_FORMATS
    
    def get_video_info(self, video_path: str) -> Dict[str, Any]:
        """获取视频文件信息"""
        try:
            probe = ffmpeg.probe(video_path)
            video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
            audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)
            
            info = {
                'file_path': video_path,
                'file_size': os.path.getsize(video_path),
                'duration': float(probe['format']['duration']),
                'format_name': probe['format']['format_name'],
            }
            
            if video_stream:
                info.update({
                    'width': int(video_stream.get('width', 0)),
                    'height': int(video_stream.get('height', 0)),
                    'fps': float(eval(video_stream.get('r_frame_rate', '0/1'))),
                    'video_codec': video_stream.get('codec_name', 'unknown')
                })
            
            if audio_stream:
                info.update({
                    'audio_codec': audio_stream.get('codec_name', 'unknown'),
                    'sample_rate': int(audio_stream.get('sample_rate', 0)),
                    'channels': int(audio_stream.get('channels', 0))
                })
            
            return info
            
        except Exception as e:
            self.logger.error(f"获取视频信息失败: {e}")
            return {}
    
    def extract_audio(self, video_path: str, output_path: Optional[str] = None, 
                     audio_format: str = 'wav') -> str:
        """从视频中提取音频"""
        try:
            if output_path is None:
                temp_dir = Path.cwd() / "temp"
                temp_dir.mkdir(exist_ok=True)
                video_name = Path(video_path).stem
                output_path = temp_dir / f"{video_name}_audio.{audio_format}"

            video = VideoFileClip(video_path)
            audio = video.audio
            
            if audio is None:
                raise ValueError("视频文件中没有音频轨道")
            
            # 为 Whisper 指定 WAV 文件的特定编解码器
            codec = 'pcm_s16le' if audio_format == 'wav' else None
            audio.write_audiofile(str(output_path), codec=codec, verbose=False, logger=None)
            
            audio.close()
            video.close()
            
            self.logger.info(f"音频提取成功: {output_path}")
            return str(output_path)
            
        except Exception as e:
            self.logger.error(f"音频提取失败: {e}")
            raise
    
    def convert_audio_format(self, input_path: str, output_path: str, 
                           target_format: str = 'wav', sample_rate: int = 16000) -> str:
        """转换音频格式"""
        try:
            audio = AudioFileClip(input_path)
            # 为 Whisper 指定 WAV 文件的特定编解码器
            codec = 'pcm_s16le' if target_format == 'wav' else None
            audio.write_audiofile(output_path, fps=sample_rate, nchannels=1, codec=codec, verbose=False, logger=None)
            audio.close()
            self.logger.info(f"音频格式转换成功: {output_path}")
            return output_path
            
        except Exception as e:
            self.logger.error(f"音频格式转换失败: {e}")
            raise
    
    def split_audio(self, audio_path: str, chunk_duration: int = 600) -> list:
        """将长音频分割成小段"""
        try:
            audio = AudioFileClip(audio_path)
            duration = audio.duration
            chunks = []
            temp_dir = Path(audio_path).parent

            for i in range(0, math.ceil(duration), chunk_duration):
                start_time = i
                end_time = min(i + chunk_duration, duration)
                chunk = audio.subclip(start_time, end_time)
                
                chunk_path = temp_dir / f"{Path(audio_path).stem}_chunk_{i//chunk_duration + 1}.wav"
                
                # 为 Whisper 指定 WAV 文件的特定编解码器
                chunk.write_audiofile(str(chunk_path), codec='pcm_s16le', verbose=False, logger=None)
                chunks.append(str(chunk_path))
                chunk.close()

            audio.close()
            self.logger.info(f"音频分割完成，共 {len(chunks)} 段")
            return chunks
            
        except Exception as e:
            self.logger.error(f"音频分割失败: {e}")
            raise
    
    def cleanup_temp_files(self, file_paths: list):
        """清理临时文件"""
        for file_path in file_paths:
            try:
                if os.path.exists(file_path):
                    os.remove(file_path)
                    self.logger.debug(f"已删除临时文件: {file_path}")
            except Exception as e:
                self.logger.warning(f"删除临时文件失败 {file_path}: {e}")